package week_07;

import java.util.HashMap;
import java.util.Map;

public class Trie {

    public class TrieNode {
        int count;
        Map<Character, TrieNode> child;

        TrieNode() {
            this.count = 0;
            this.child = new HashMap<>();
        }
    }

    TrieNode root;

    public Trie() {
        root = new TrieNode();
    }

    /**
     * Inserts a word into the trie.
     */
    public void insert(String word) {
        solve(word, true, false);
    }

    /**
     * Returns if the word is in the trie.
     */
    public boolean search(String word) {
        return solve(word, false, false);
    }

    /**
     * Returns if there is any word in the trie that starts with the given prefix.
     */
    public boolean startsWith(String prefix) {
        return solve(prefix, false, true);
    }

    public boolean solve(String word, boolean isInsert, boolean isPrefix) {
        TrieNode cur = root;
        for (char c : word.toCharArray()) {
            if (!cur.child.containsKey(c)) {
                if (isInsert) {
                    cur.child.put(c, new TrieNode());
                } else {
                    return false;
                }
            }
            cur = cur.child.get(c);
        }
        if (isInsert) cur.count += 1;
        if (isPrefix) return true;
        return cur.count > 0;
    }
}
